from copy import copy
from typing import Dict, List, Optional, Tuple, Union


class LMTemplateParser:
    """
    语言模型的提示模板解析器。
    用于处理和格式化模型的提示模板。
    
    Args:
        meta_template: 模型的元模板配置列表
    """

    def __init__(self, meta_template: Optional[List[Dict]] = None):
        self.meta_template = meta_template
        if meta_template:
            assert isinstance(meta_template, list)
            # 存储角色名到配置的映射
            self.roles: Dict[str, dict] = dict()
            for item in meta_template:
                assert isinstance(item, dict)
                assert item['role'] not in self.roles, '角色名必须唯一!'
                self.roles[item['role']] = item.copy()

    def __call__(self, dialog) -> str:
        """
        解析提示模板，并在需要时用元模板包装。

        Args:
            dialog: 提示模板（可能需要被元模板包装）

        Returns:
            str: 最终的提示字符串
        """
        assert isinstance(dialog, (str, list))
        if isinstance(dialog, str):
            return dialog
            
        if self.meta_template:
            # 使用元模板处理
            prompt = ''
            for index, item in enumerate(dialog):
                if isinstance(item, str):
                    prompt += item
                else:
                    new_str = self._prompt2str(item, index == len(dialog) - 1)
                    prompt += new_str
        else:
            # 不使用元模板的简单处理
            prompt = ''
            last_sep = ''
            for item in dialog:
                if isinstance(item, str):
                    if item:
                        prompt += last_sep + item
                elif item.get('content', ''):
                    prompt += last_sep + item.get('prompt', '')
                last_sep = '\n'
        return prompt

    def _format_begin(self, role_cfg, message):
        """格式化角色开始标记"""
        name = message.get('name', None)
        if name is not None:
            begin = role_cfg['begin'].get('with_name', '')
            if name in role_cfg['begin'].get('name', {}):
                begin = begin.format(name=role_cfg['begin']['name'][name])
            else:
                begin = begin.format(name=name)
        else:
            if isinstance(role_cfg.get('begin', ''), str):
                begin = role_cfg.get('begin', '')
            elif isinstance(role_cfg['begin'], dict):
                begin = role_cfg['begin'].get('without_name', '')
        return begin

    def _prompt2str(self, prompt: Union[str, Dict], last: bool = False) -> Tuple[str, bool]:
        """
        将提示转换为字符串。

        Args:
            prompt: 提示信息，可以是字符串或字典。
            last: 布尔值，指示是否为最后一个提示。

        Returns:
            Tuple[str, bool]: 转换后的字符串和布尔值。
        """
        # 如果提示是字符串，直接返回
        if isinstance(prompt, str):
            return prompt

        # 获取角色配置，如果有备用角色则使用备用角色
        merged_prompt = self.roles.get(prompt['role'])
        if merged_prompt.get('fallback_role'):
            merged_prompt = self.roles.get(merged_prompt['fallback_role'])

        # 格式化角色开始标记
        begin = self._format_begin(merged_prompt, prompt)
        res = begin

        # 如果是最后一个提示且需要生成内容，直接返回内容
        if last and merged_prompt.get('generate', False):
            res += prompt.get('content', '')
            return res

        # 添加提示内容和结束标记
        res += prompt.get('content', '') + merged_prompt.get('end', '')

        # 如果是最后一个提示且角色不是助手，添加助手的开始标记
        if last and merged_prompt['role'] != 'assistant':
            res += self._format_begin(self.roles['assistant'], {})

        return res


class BaseLLM:
    """
    语言模型包装器的基类。
    
    Args:
        path: 模型路径
        max_new_tokens: 模型生成的最大token数，默认512
        tokenizer_only: 是否只初始化分词器，默认False
        meta_template: 模型的元提示模板（如果需要）
    """

    def __init__(self,
                 path: str,
                 tokenizer_only: bool = False,
                 template_parser: 'LMTemplateParser' = LMTemplateParser,
                 meta_template: Optional[List[Dict]] = None,
                 *,
                 max_new_tokens: int = 512,
                 top_p: float = 0.8,
                 top_k: float = 40,
                 temperature: float = 0.8,
                 repetition_penalty: float = 1.0,
                 stop_words: Union[List[str], str] = None):
        """初始化模型参数"""
        self.path = path
        self.tokenizer_only = tokenizer_only
        self.template_parser = template_parser(meta_template)
        
        # 设置结束标记ID
        self.eos_token_id = None
        if meta_template and 'eos_token_id' in meta_template:
            self.eos_token_id = meta_template['eos_token_id']

        # 处理停止词
        if isinstance(stop_words, str):
            stop_words = [stop_words]
            
        # 生成参数配置
        self.gen_params = dict(
            max_new_tokens=max_new_tokens,  # 最大生成长度
            top_p=top_p,  # 累积概率阈值
            top_k=top_k,  # 保留最高概率的k个token
            temperature=temperature,  # 采样温度
            repetition_penalty=repetition_penalty,  # 重复惩罚系数
            stop_words=stop_words  # 停止词列表
        )

    # 以下是需要子类实现的接口方法
    def generate(self, inputs: Union[str, List[str]], **gen_params) -> str:
        """生成文本"""
        raise NotImplementedError

    def stream_generate(self, inputs: str, **gen_params) -> List[str]:
        """流式生成文本"""
        raise NotImplementedError

    def chat(self, inputs: Union[List[dict], List[List[dict]]], 
             session_ids: Union[int, List[int]] = None,
             **gen_params):
        """对话生成"""
        if isinstance(inputs[0], list):
            _inputs = list()
            for msg in inputs:
                _inputs.append(self.template_parser(msg))
        else:
            _inputs = self.template_parser(inputs)
        return self.generate(_inputs, **gen_params)

    def stream_chat(self, inputs: List[dict], **gen_params):
        """流式对话"""
        raise NotImplementedError

    def tokenize(self, prompts):
        """分词处理"""
        raise NotImplementedError

    def update_gen_params(self, **kwargs):
        """更新生成参数"""
        gen_params = copy(self.gen_params)
        gen_params.update(kwargs)
        return gen_params


# 异步支持
class AsyncLLMMixin:
    """异步语言模型混入类，提供异步接口"""
    
    # 异步版本的接口方法
    async def generate(self, inputs, session_ids=None, **gen_params):
        raise NotImplementedError

    async def stream_generate(self, inputs: str, **gen_params):
        raise NotImplementedError

    async def chat(self, inputs, session_ids=None, **gen_params):
        if isinstance(inputs[0], list):
            _inputs = list()
            for msg in inputs:
                _inputs.append(self.template_parser(msg))
        else:
            _inputs = self.template_parser(inputs)
        return await self.generate(_inputs, session_ids, **gen_params)

    async def stream_chat(self, inputs: List[dict], **gen_params):
        raise NotImplementedError

    async def tokenize(self, prompts):
        raise NotImplementedError


class AsyncBaseLLM(AsyncLLMMixin, BaseLLM):
    """异步语言模型基类，结合异步接口和基础功能"""
    pass
